package com.xinan.ollama.api.utils;


import com.alibaba.fastjson.JSON;
import com.xinan.ollama.api.entity.OllamaRequstObj;
import com.xinan.ollama.api.vo.req.OllamaReq;
import com.xinan.ollama.api.vo.res.ChatRes;
import com.xinan.ollama.api.vo.res.EmbeddingRes;
import com.xinan.ollama.api.vo.res.GenerateRes;
import com.xinan.ollama.config.OkHttpClientConfig;
import com.xinan.ollama.config.OllamaConfig;
import lombok.extern.slf4j.Slf4j;
import okhttp3.*;
import org.springframework.stereotype.Component;

import javax.annotation.PostConstruct;
import javax.annotation.Resource;
import java.io.IOException;
import java.util.HashMap;
import java.util.Map;

/**
 * Ollama接口调用
 *
 * @author yuanz
 * @date 2025-02-21
 */
@Slf4j
@Component
public class OllamaUtil {

    @Resource
    private OllamaConfig ollamaConfig;

    private static final int CODE200 = 200;

    /**
     * 进行模型的首次调用
     */
    @PostConstruct
    private void init() {
        String model = ollamaConfig.getModel();
        OllamaReq ollamaReq = new OllamaReq();
        ollamaReq.setModel(model);
        ollamaReq.setPrompt("你是谁?");
        long start = System.currentTimeMillis();
        GenerateRes generate = generate(ollamaReq);
        log.info("Init [{}] result:[{}];Cost:[{}]ms.", model, generate.getResponse(), System.currentTimeMillis() - start);
    }

    public GenerateRes generate(OllamaReq req) {
        try {
            OllamaRequstObj ollamaRequstObj = new OllamaRequstObj();
            ollamaRequstObj.setModel(req.getModel());
            ollamaRequstObj.setPrompt(req.getPrompt());
            String bodyStr = ollamaRemoteApi("/api/generate", ollamaRequstObj);
            return JSON.parseObject(bodyStr, GenerateRes.class);
        } catch (IOException e) {
            e.printStackTrace();
        }
        return new GenerateRes();
    }

    public ChatRes chat(OllamaRequstObj req) {
        try {
            String bodyStr = ollamaRemoteApi("/api/chat", req);
            return JSON.parseObject(bodyStr, ChatRes.class);
        } catch (IOException e) {
            e.printStackTrace();
        }
        return new ChatRes();
    }

    public Map<String, Object> getEventAndOrganizeNameByInformationName(OllamaReq req) {
        try {
            String bodyStr = ollamaRemoteApi("/api/chat", req);
            return JSON.parseObject(bodyStr, Map.class);
        } catch (IOException e) {
            e.printStackTrace();
        }
        return new HashMap<>();
    }

    public EmbeddingRes getEmbeddingByStr(OllamaReq req) {
        try {
            String bodyStr = ollamaRemoteApi("/api/embeddings", req);
            return JSON.parseObject(bodyStr, EmbeddingRes.class);
        } catch (IOException e) {
            e.printStackTrace();
        }
        return new EmbeddingRes();
    }

    private String ollamaRemoteApi(String url, Object req) throws IOException {
        String bodyStr = "";
        OkHttpClient client = OkHttpClientConfig.getUnsafeOkHttpClient();
        MediaType mediaType = MediaType.parse("application/json");
        RequestBody body = RequestBody.create(mediaType, JSON.toJSONString(req));
        Request request = new Request.Builder()
                .url(ollamaConfig.baseUrl + url)
                .method("POST", body)
                .addHeader("User-Agent", "Apifox/1.0.0 (https://apifox.com)")
                .addHeader("Content-Type", "application/json")
                .addHeader("Accept", "*/*")
                .addHeader("Connection", "keep-alive")
                .build();
        Response response = client.newCall(request).execute();
        if (response.code() == CODE200) {
            bodyStr = response.body().string();
        }
        return bodyStr;
    }

}
